"""
基于Transformer的多资产时序预测策略
策略特点：
1. 使用多头注意力捕捉跨资产关联
2. 自适应时间窗口处理不同频率数据
3. 集成风险控制模块
"""

import torch
import torch.nn as nn
import numpy as np
from torch.nn import TransformerEncoder, TransformerEncoderLayer

class MultiAssetTransformer(nn.Module):
    def __init__(self, n_assets=5, d_model=64, nhead=4):
        super().__init__()
        self.asset_embedding = nn.Embedding(n_assets, d_model)
        self.time_embedding = nn.Linear(1, d_model)
        
        encoder_layers = TransformerEncoderLayer(
            d_model, nhead, dim_feedforward=256)
        self.transformer = TransformerEncoder(encoder_layers, 3)
        
        self.pred_head = nn.Sequential(
            nn.Linear(d_model, 128),
            nn.ReLU(),
            nn.Linear(128, n_assets)
        )
        
        self.risk_module = nn.Sequential(
            nn.Linear(d_model*n_assets, 64),
            nn.ReLU(),
            nn.Linear(64, n_assets)
        )
        
    def forward(self, x, timesteps):
        # x: [batch, seq_len, n_assets]
        # timesteps: [batch, seq_len]
        batch_size, seq_len = x.shape[0], x.shape[1]
        
        # 资产嵌入
        asset_ids = torch.arange(x.shape[2]).to(x.device)
        asset_emb = self.asset_embedding(asset_ids)  # [n_assets, d_model]
        
        # 时间嵌入
        time_emb = self.time_embedding(
            timesteps.float().unsqueeze(-1))  # [batch, seq_len, d_model]
        
        # 组合特征
        x = x.permute(2, 0, 1)  # [n_assets, batch, seq_len]
        x = asset_emb.unsqueeze(1).unsqueeze(1) + time_emb.permute(1,0,2)
        x = x.reshape(n_assets*batch_size, seq_len, -1)
        
        # Transformer处理
        x = self.transformer(x)
        x = x.mean(dim=1)  # 聚合时序特征
        
        # 预测和风险
        pred = self.pred_head(x)
        risk = self.risk_module(x.reshape(batch_size, -1))
        
        return pred, risk

class MultiAssetStrategy:
    def __init__(self, n_assets=5):
        self.model = MultiAssetTransformer(n_assets)
        self.optimizer = torch.optim.Adam(
            self.model.parameters(), lr=0.001)
        self.loss_fn = nn.MSELoss()
        
    def train(self, dataloader, epochs=10):
        self.model.train()
        for epoch in range(epochs):
            for x, t, y in dataloader:
                self.optimizer.zero_grad()
                pred, risk = self.model(x, t)
                loss = self.loss_fn(pred, y) + 0.1*risk.mean()
                loss.backward()
                self.optimizer.step()
                
    def predict(self, x, t):
        self.model.eval()
        with torch.no_grad():
            return self.model(x, t)
            
    def get_allocations(self, market_data, timestamps):
        pred, risk = self.predict(market_data, timestamps)
        return torch.softmax(pred - risk, dim=-1)

if __name__ == '__main__':
    # 示例用法
    strategy = MultiAssetStrategy()
    
    # 模拟数据
    batch_size = 32
    seq_len = 60
    n_assets = 5
    x = torch.randn(batch_size, seq_len, n_assets)
    t = torch.arange(seq_len).repeat(batch_size, 1)
    y = torch.randn(batch_size, n_assets)
    
    # 训练
    class Dataset(torch.utils.data.Dataset):
        def __init__(self, x, t, y):
            self.x = x
            self.t = t
            self.y = y
            
        def __len__(self):
            return len(self.x)
            
        def __getitem__(self, idx):
            return self.x[idx], self.t[idx], self.y[idx]
    
    dataloader = torch.utils.data.DataLoader(
        Dataset(x, t, y), batch_size=8)
    
    strategy.train(dataloader)
    
    # 预测
    test_x = torch.randn(1, seq_len, n_assets)
    test_t = torch.arange(seq_len).reshape(1, -1)
    alloc = strategy.get_allocations(test_x, test_t)
    print("资产配置权重:", alloc.detach().numpy())